# The function name, type hints, and docstring are all part of the tool
# schema that's passed to the model. Defining good, descriptive schemas
# is an extension of prompt engineering and is an important part of
# getting models to perform well.
# 1. 这是第一种工具的方式  使用python函数

# def add(a: int, b: int) -> int:
#     """Add two integers.
#
#     Args:
#         a: First integer
#         b: Second integer
#     """
#     return a + b
#
#
# def multiply(a: int, b: int) -> int:
#     """Multiply two integers.
#
#     Args:
#         a: First integer
#         b: Second integer
#     """
#     return a * b

# 2. 这是第二种工具的方式  使用pydantic class

# from pydantic import BaseModel,Field
#
# class add(BaseModel):
#     a:int=Field(...,description="第一个整数")
#     b:int=Field(...,description="第二个整数")
#
# class multiply(BaseModel):
#     a:int=Field(...,description="第一个整数")
#     b:int=Field(...,description="第二个整数")

# 3. 这是第三种工具的方式  使用typedict class
# from typing_extensions import Annotated, TypedDict
#
# class add(TypedDict):
#     """Add two integers."""
#
#     # Annotations must have the type and can optionally include a default value and description (in that order).
#     a: Annotated[int, ..., "First integer"]
#     b: Annotated[int, ..., "Second integer"]
#
# class multiply(TypedDict):
#     """Multiply two integers."""
#
#     a: Annotated[int, ..., "First integer"]
#     b: Annotated[int, ..., "Second integer"]
#

# 4. 这是第四种工具的方式  使用langchain tools
from langchain_core.tools import tool

# 加上了一个修饰器
@tool
def add(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b
@tool
def multiply(a: int, b: int) -> int:
    """Multiply two numbers."""
    return a * b

print(multiply.name)
print(multiply.description)
print(multiply.args)

tools = [add, multiply]

# 绑定tool
from langchain.chat_models import init_chat_model
import os

key = os.getenv("OPENAI_API_KEY")
# print(key)
api_key = str(key)

llm = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    base_url="https://api.zetatechs.com/v1",
    api_key=api_key
)
llm_with_tools = llm.bind_tools(tools)  # 这里tools 是一个工具列表
query = "What is 3 * 12?"

print(llm_with_tools.invoke(query))
query = "What is 3 * 12? Also, what is 11 + 49?"
result = llm_with_tools.invoke(query)
print(result)
print(llm_with_tools.invoke(query).tool_calls)


# 结构解析器 Pydantic
from langchain_core.output_parsers import PydanticToolsParser
from pydantic import BaseModel, Field


class add(BaseModel):
    """Add two integers."""

    a: int = Field(..., description="First integer")
    b: int = Field(..., description="Second integer")


class multiply(BaseModel):
    """Multiply two integers."""

    a: int = Field(..., description="First integer")
    b: int = Field(..., description="Second integer")


chain = llm_with_tools | PydanticToolsParser(tools=[add, multiply])
result = chain.invoke(query)
print(result)
for chain in result:
    print(chain)